# -*- coding: utf-8 -*-  
'''
数据集

@author: luoyi
Created on 2021年3月17日
'''
import tensorflow as tf

import utils.conf as conf
import data.dataset as ds

#    数据迭代器（含数据预处理工作）
def tensor_data_iterator(in_path, 
                         out_path,
                         count=10000,
                         sentence_maxlen=conf.TRANSFORMER.get_sentence_maxlen(),
                         x_preprocess=None,
                         y_preprocess=None):
    couplet_iterator = ds.read_couplet_iterator(in_path, out_path, count)
    for in_line, out_line in couplet_iterator:
        X = x_preprocess(frist_couplet=in_line, second_couplet=out_line, sentence_maxlen=sentence_maxlen)
        Y = y_preprocess(second_couplet=out_line, sentence_maxlen=sentence_maxlen)
        
        yield X, Y
        pass
    pass


#    tensor数据集
def tensor_db(in_path,
              out_path,
              count,
              sentence_maxlen=conf.TRANSFORMER.get_sentence_maxlen() + 1,
              x_preprocess=None,
              y_preprocess=None,
              batch_size=conf.DATASET.get_batch_size(),
              epochs=conf.DATASET.get_epochs(),
              shuffle_buffer_rate=conf.DATASET.get_shuffle_buffer_rate()):
    x_shape = tf.TensorShape([2, sentence_maxlen])
    y_shape = tf.TensorShape((sentence_maxlen, ))
    db = tf.data.Dataset.from_generator(generator=lambda :tensor_data_iterator(in_path=in_path,
                                                                               out_path=out_path,
                                                                               count=count,
                                                                               sentence_maxlen=sentence_maxlen,
                                                                               x_preprocess=x_preprocess,
                                                                               y_preprocess=y_preprocess), 
                                        output_types=(tf.float32, tf.float32), 
                                        output_shapes=(x_shape, y_shape))
    if (shuffle_buffer_rate > 0): db = db.shuffle(buffer_size=batch_size * shuffle_buffer_rate)
    if (batch_size > 0): db = db.batch(batch_size=batch_size, drop_remainder=True)
    if (epochs > 0): db = db.repeat(epochs)
    return db


